from config import *
import torch
from torch.utils.data import Dataset
import numpy as np
import copy
from easySample import easySample
from dataPre import dataPre
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")

# load_from_shelve=True
load_from_shelve = False


class RGCNDataset(Dataset):
    def __init__(self, datas, datasName, is_train):
        self.data = datas
        self.samples = {}
        self.eSample = easySample()
        self.is_train = is_train
        self.datasName = datasName
        self.tripletDatas = []

        if self.is_train:
            lth = len(datas)//2
            for i in range(lth):
                if datas[i][0] != datas[lth+i][0]:
                    print("error 样本对不匹配")
                self.tripletDatas.append(
                    (datas[i][0], datas[i][1], datas[lth+i][1]))

    def sampels_trans_to_batch(self, lth, att_dimension, batch, ind):
        max_node_size = -float("inf")
        for i in range(lth):
            max_node_size = max(max_node_size, len(batch[i][ind][0]))

        batch_adj, batch_att, batch_vtype, batch_adj_weight = [], [], [], []
        for i in range(lth):
            temp_adj = np.zeros((max_node_size, max_node_size))
            temp_att = np.zeros((max_node_size, att_dimension))
            temp_vtype = np.zeros((max_node_size, 3))

            sample_size = len(batch[i][ind][0])  # 当前样本的长度

            temp_adj[:sample_size, :sample_size] = batch[i][ind][0]
            temp_att[:sample_size, :] = batch[i][ind][1]
            temp_vtype[:sample_size, :] = batch[i][ind][2]

            batch_adj.append(copy.deepcopy(temp_adj))  # 深拷贝,免得出现错误
            batch_att.append(copy.deepcopy(temp_att))
            batch_vtype.append(copy.deepcopy(temp_vtype))

        res = [torch.tensor(np.array(batch_adj)),
               torch.tensor(np.array(batch_att)),
               torch.tensor(np.array(batch_vtype))]
        return res

    def adjust_samples_to_same_dimension(self, batch):
        lth = len(batch)
        att_dimension = len(batch[0][0][1][0])

        batch_value = []
        res_x, res_y, res_z = [], [], []
        res_x = self.sampels_trans_to_batch(lth, att_dimension, batch, 0)
        res_y = self.sampels_trans_to_batch(lth, att_dimension, batch, 1)
        if self.is_train:
            res_z = self.sampels_trans_to_batch(lth, att_dimension, batch, 2)
        else:
            for i in range(lth):
                batch_value.append(batch[i][3][0])
        return (torch.tensor(batch_value), res_x, res_y, res_z)

    def __getitem__(self, idx):
        if self.is_train:
            sample_x = self.eSample.get_sample(
                self.tripletDatas[idx][0], modelName)
            sample_y = self.eSample.get_sample(
                self.tripletDatas[idx][1], modelName)
            sample_z = self.eSample.get_sample(
                self.tripletDatas[idx][2], modelName)
            x = [sample_x["adj"], sample_x["att"], sample_x["vtype"]]
            y = [sample_y["adj"], sample_y["att"], sample_y["vtype"]]
            z = [sample_z["adj"], sample_z["att"], sample_z["vtype"]]
            return [x, y, z, [0]]
        else:
            sample_x = self.eSample.get_sample(self.data[idx][0], modelName)
            sample_y = self.eSample.get_sample(self.data[idx][1], modelName)
            x = [sample_x["adj"], sample_x["att"], sample_x["vtype"]]
            y = [sample_y["adj"], sample_y["att"], sample_y["vtype"]]
            value = torch.tensor([self.data[idx][2]])
            return [x, y, [0], value]

    def __len__(self):
        if self.is_train:
            return len(self.tripletDatas)
        return len(self.data)


if __name__ == "__main__":
    datapre = dataPre()
    data_pair_infs = datapre.get_pair_infs()
    load_from_shelve = False
    print("处理RGCN数据集的预加载")
    dataset = RGCNDataset(data_pair_infs["test"], "test", False)
    dataset = RGCNDataset(data_pair_infs["valid"], "valid", False)
    dataset = RGCNDataset(data_pair_infs["train"], "train", False)
    print("RGCN数据集的预加载完成")
